import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.animation import FuncAnimation, FFMpegWriter
import os
from scipy.integrate import odeint


class Animate(object):
    """
    动画
    """
    def __init__(self):
        self.fig, self.ax = plt.subplots()

        ffmpegpath = os.path.abspath("D:\\software\\ffmpeg-latest-win64-static\\bin\\ffmpeg.exe")
        mpl.rcParams["animation.ffmpeg_path"] = ffmpegpath
        self.writer = FFMpegWriter()

    def sin_1(self):
        """
        sin(x)动画演示
        :return:
        """
        x, y = [], []
        line, = self.ax.plot([], [], 'r-', animated=False)  # 动画画图

        def init():
            self.ax.set_xlim(0, 2 * np.pi)
            self.ax.set_ylim(-1, 1)
            return line,

        def update(frame):
            x.append(frame)
            y.append(np.sin(frame))
            line.set_data(x, y)
            return line,

        anim = FuncAnimation(self.fig, update, frames=np.linspace(0, 2 * np.pi, 128),
                             init_func=init, blit=True, interval=30)

        # anim.save('./data/sin_1.mp4', writer=self.writer)
        anim.save('./data/sin_1.gif', writer='imagemagick', fps=30)

    def sin_2(self):
        """
        sin(x)动画演示
        :return:
        """
        x = np.linspace(0, 2 * np.pi, 200)
        y = np.sin(x)
        l = self.ax.plot(x, y)  # 背景图
        dot, = self.ax.plot([], [], 'ro')  # 动画画图
        x_dot, y_dot = [], []

        def init():
            self.ax.set_xlim(0, 2 * np.pi)
            self.ax.set_ylim(-1, 1)
            return l,

        def update_dot(frame):
            dot.set_data(frame, np.sin(frame))
            return dot,

        anim = FuncAnimation(self.fig, update_dot, frames=x, interval=100, init_func=init)
        anim.save('./data/sin_2.gif', writer='imagemagick', fps=30)

        plt.show()

    def simple_pendulum(self):
        """
        单摆动画
        :return:
        """
        g = 9.8
        leng = 2.0

        def pendulum_equations1(w, t, l):
            th, dth = w
            dv = -g * np.sin(th) / l
            return dth, dv

        def pendulum_equation2(w, t, l, b):
            th, dth = w
            dv = (-b * dth - g * np.sin(th)) / l
            return dth, dv

        t = np.arange(0, 20, 0.1)
        track = odeint(pendulum_equations1, (1.0, 0), t, args=(leng,))

        x = [leng * np.sin(track[i, 0]) for i in range(len(track))]
        y = [-leng * np.cos(track[i, 0]) for i in range(len(track))]

        self.ax.grid()
        line, = self.ax.plot([], [], 'o-', lw=2)
        time_template = 'time = {}'
        time_text = self.ax.text(0.05, 0.9, '', transform=self.ax.transAxes)

        def init():
            self.ax.set_xlim(-3, 3)
            self.ax.set_ylim(-3, 3)
            time_text.set_text('')
            return line, time_text

        def update(i):
            new_x = [0, x[i]]
            new_y = [0, y[i]]
            line.set_data(new_x, new_y)
            time_text.set_text(time_template.format(0.1 * i))
            return line, time_text

        anim = FuncAnimation(self.fig, update, frames=range(1, len(x)), init_func=init,
                             interval=50)

        anim.save('./data/simple_pendulum.gif', writer='imagemagick', fps=30)
        plt.show()



animation = Animate()
# animation.sin_1()
# animation.sin_2()
animation.simple_pendulum()

